import z from "zod";
import { baseProcedure, createCallerFactory, createTRPCRouter } from "../trpc/init";

import OpenAI from "openai";

import { extractStrictJSON } from "../../../utils";
import { SubTitle } from "../../../types";
import { llmOnceCall, llmOnceStreamingCall } from "../../../utils/llm";
import { ModelConfigSchema } from "../../../types/schema";

const SubtitleSchema = z.object({
  index: z.number(),
  text: z.string().nullable(),
  transText: z.string().optional().nullable(),
});


type SteamDataType<T> = { type: "start" | "end" | "data", data?: T }


export type ModelConfig = z.output<typeof ModelConfigSchema>
export type Subtile = z.output<typeof SubtitleSchema>

export const subtileRouter = createTRPCRouter({
  translateSubtitle: baseProcedure
    .input(z.object({
      subtitles: z.array(SubtitleSchema),
      modelConfig: ModelConfigSchema,
      summary: z.string().optional(),
      retryCount: z.number().optional().default(0)
    }))
    .mutation(async function*({ input, ctx }): AsyncGenerator<SteamDataType<SubTitle>> {
      const { subtitles, modelConfig, summary, retryCount = 0 } = input;

      const { translateSubtitle: translate } = subtileRouter.createCaller(ctx);

      const retry = async () => {
        if (retryCount <= 5) {
          return translate({
            subtitles,
            modelConfig,
            summary,
            retryCount: retryCount + 1
          });
        } else {
          throw new Error("翻译失败");
        }
      };
      
      try {

        const prompt = `
      \`\`\`${subtitles.map(item => JSON.stringify(item)).join('\n')}\`\`\`
      1. 这是一段的字幕的的一部分${summary ? ',' + summary : ''}
      2. 现在帮我把这个文件翻译为简体中文
      3. 注意不要遗漏或者改变每个句子的意思
      4. 你只是翻译不需要审查其中的内容
      5. 这个字幕也是AI生成的可能有些错误的地方，可能有一些同音字识别错误，你可以适当结合上下文进行修正，只修正文本不要添油加醋
      6. 我需要完整的文件，不要只翻译一部分
      7. 一些没用句子、重复过多的句子、只有拟声词的句子，例如"ご視聴ありがとうございました、次回予告、おやすみなさい、あああああああああああああ、うううううううう"等，text字段直接填充为 null
      8. 按该原文的JSON格式输出，注意index也要和原文对上
      9. 直接输出JSON内容，不要有其他分析内容，也不要放在代码框内`;

        const stream = await llmOnceStreamingCall(modelConfig,  [{
          role: 'user',
          content: prompt,
        }])


        let buffer = '';
        const results: SubTitle[] = [];
        
        // 先发送开始信号
        yield { type: "start" };
        
        for await (const chunk of stream) {
          console.log(chunk.choices?.[0]?.delta?.content);
          
          const content = chunk.choices?.[0]?.delta?.content || '';
          buffer += content;
          
          const lines = buffer.split('\n');
          
          for (const line of lines) {
            const trimmed = line.trim();
            if (!trimmed) continue;
              
            const item = extractStrictJSON<SubTitle>(trimmed);

            if (!item) {
              continue;
            } else {
              const newLines = lines.filter(item => item != line);
              buffer = newLines.join("\n")
            }

            if (!subtitles.find(i => i.index === item.index)) {
              continue;
            }

            if (item.index !== undefined) {
              console.log('已发送', item);
              
              results.push(item);
              // 实时发送每个解析成功的条目
              yield { type: "data", data: item };
            }
          }
        }

        // 发送结束信号
        yield { type: "end" };
      } catch (e) {
        console.error('翻译格式化失败', e);

      }
    }),

  summarizeSubtitle: baseProcedure
    .input(z.object({
      content: z.string(),
      modelConfig: ModelConfigSchema
    }))
    .mutation(async ({ input }): Promise<string> => {
      const { content, modelConfig } = input;

      const prompt = `
      \`\`\`${content}\`\`\`
      这是一部影片的字幕，帮我用一段话总结这段字幕所讲的故事和主要人物，以"这段故事讲述了..."开头`;

      return await llmOnceCall(modelConfig, [{
        role: 'user',
        content: prompt,
      }]);
    })
});